import cv2
import numpy as np

# 数据增强函数
def augment_image(image):
    # 确保输入图像为灰度图（单通道）
    image = image.astype(np.uint8)

    # 随机旋转
    angle = np.random.uniform(-50, 50)  # 随机旋转角度
    M = cv2.getRotationMatrix2D((14, 14), angle, 1.0)  # 计算旋转矩阵
    rotated = cv2.warpAffine(image, M, (28, 28))  # 进行旋转

    # 随机平移
    tx = np.random.randint(-3, 3)  # 随机水平平移
    ty = np.random.randint(-3, 3)  # 随机垂直平移
    translated = cv2.warpAffine(rotated, np.float32([[1, 0, tx], [0, 1, ty]]), (28, 28))

    # 随机缩放
    scale = np.random.uniform(0.9, 1.1)  # 随机缩放因子
    new_size = int(28 * scale)  # 计算缩放后的尺寸
    scaled_image = cv2.resize(translated, (new_size, new_size))  # 缩放图像

    # 如果缩放后尺寸小于28x28，填充图像到28x28
    if new_size < 28:
        delta_w = 28 - new_size
        delta_h = 28 - new_size
        top, bottom = delta_h // 2, delta_h - (delta_h // 2)
        left, right = delta_w // 2, delta_w - (delta_w // 2)
        scaled = cv2.copyMakeBorder(scaled_image, top, bottom, left, right, cv2.BORDER_CONSTANT, value=0)
    else:
        # 如果缩放后尺寸大于28x28，裁剪图像到28x28
        start_x = (new_size - 28) // 2
        start_y = (new_size - 28) // 2
        scaled = scaled_image[start_y:start_y + 28, start_x:start_x + 28]

    return scaled